#!/usr/bin/env python3
"""
API安全中间件
提供请求频率限制、输入验证、安全头等功能
"""

import time
import hashlib
import re
import json
from functools import wraps
from flask import request, jsonify, g
from datetime import datetime, timedelta
import redis
from typing import Dict, Any, Optional

class SecurityMiddleware:
    """安全中间件类"""
    
    def __init__(self, redis_client):
        self.redis_client = redis_client
        
        # 频率限制配置
        self.rate_limits = {
            'login': {'requests': 5, 'window': 300},      # 5次/5分钟
            'sms': {'requests': 3, 'window': 300},        # 3次/5分钟
            'register': {'requests': 3, 'window': 3600},  # 3次/1小时
            'default': {'requests': 100, 'window': 3600}  # 100次/1小时
        }
        
        # 敏感字段列表
        self.sensitive_fields = [
            'password', 'token', 'secret', 'key', 'auth',
            'credential', 'private', 'confidential'
        ]
        
        # SQL注入检测模式
        self.sql_injection_patterns = [
            r"(\b(SELECT|INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|EXEC|UNION)\b)",
            r"(\b(OR|AND)\s+\d+\s*=\s*\d+)",
            r"(\b(OR|AND)\s+['\"]?\w+['\"]?\s*=\s*['\"]?\w+['\"]?)",
            r"(--|#|/\*|\*/)",
            r"(\bUNION\s+SELECT\b)",
            r"(\bINTO\s+OUTFILE\b)",
            r"(\bLOAD_FILE\b)"
        ]
        
        # XSS检测模式
        self.xss_patterns = [
            r"<script[^>]*>.*?</script>",
            r"javascript:",
            r"on\w+\s*=",
            r"<iframe[^>]*>",
            r"<object[^>]*>",
            r"<embed[^>]*>",
            r"<link[^>]*>",
            r"<meta[^>]*>"
        ]

    def get_client_identifier(self) -> str:
        """获取客户端标识符"""
        # 优先使用真实IP
        real_ip = request.headers.get('X-Real-IP')
        if real_ip:
            return real_ip
        
        # 其次使用转发IP
        forwarded_for = request.headers.get('X-Forwarded-For')
        if forwarded_for:
            return forwarded_for.split(',')[0].strip()
        
        # 最后使用远程地址
        return request.remote_addr or 'unknown'

    def check_rate_limit(self, endpoint: str, identifier: str = None) -> tuple[bool, dict]:
        """检查请求频率限制"""
        if identifier is None:
            identifier = self.get_client_identifier()
        
        # 获取对应端点的限制配置
        limit_config = self.rate_limits.get(endpoint, self.rate_limits['default'])
        max_requests = limit_config['requests']
        window_seconds = limit_config['window']
        
        # Redis键
        key = f"rate_limit:{endpoint}:{identifier}"
        
        try:
            # 获取当前计数
            current_count = self.redis_client.get(key)
            
            if current_count is None:
                # 首次请求
                self.redis_client.setex(key, window_seconds, 1)
                return True, {
                    'allowed': True,
                    'remaining': max_requests - 1,
                    'reset_time': int(time.time()) + window_seconds
                }
            else:
                current_count = int(current_count)
                
                if current_count >= max_requests:
                    # 超出限制
                    ttl = self.redis_client.ttl(key)
                    return False, {
                        'allowed': False,
                        'remaining': 0,
                        'reset_time': int(time.time()) + ttl,
                        'retry_after': ttl
                    }
                else:
                    # 增加计数
                    self.redis_client.incr(key)
                    ttl = self.redis_client.ttl(key)
                    return True, {
                        'allowed': True,
                        'remaining': max_requests - current_count - 1,
                        'reset_time': int(time.time()) + ttl
                    }
                    
        except Exception as e:
            # Redis错误时允许请求通过，但记录错误
            print(f"Rate limit check failed: {str(e)}")
            return True, {'allowed': True, 'error': str(e)}

    def validate_input_data(self, data: Any, max_depth: int = 5) -> tuple[bool, str]:
        """深度验证输入数据"""
        if max_depth <= 0:
            return False, "数据嵌套层级过深"
        
        if isinstance(data, dict):
            for key, value in data.items():
                # 检查键名
                if not self._is_safe_string(str(key)):
                    return False, f"不安全的键名: {key}"
                
                # 递归检查值
                is_valid, message = self.validate_input_data(value, max_depth - 1)
                if not is_valid:
                    return False, message
                    
        elif isinstance(data, list):
            if len(data) > 1000:  # 限制数组长度
                return False, "数组长度超出限制"
            
            for item in data:
                is_valid, message = self.validate_input_data(item, max_depth - 1)
                if not is_valid:
                    return False, message
                    
        elif isinstance(data, str):
            if len(data) > 10000:  # 限制字符串长度
                return False, "字符串长度超出限制"
            
            if not self._is_safe_string(data):
                return False, "检测到不安全的字符串内容"
                
        return True, "验证通过"

    def _is_safe_string(self, text: str) -> bool:
        """检查字符串是否安全"""
        # SQL注入检测
        for pattern in self.sql_injection_patterns:
            if re.search(pattern, text, re.IGNORECASE):
                return False
        
        # XSS检测
        for pattern in self.xss_patterns:
            if re.search(pattern, text, re.IGNORECASE):
                return False
        
        # 检查是否包含null字节
        if '\x00' in text:
            return False
        
        return True

    def sanitize_data(self, data: Any) -> Any:
        """清理数据"""
        if isinstance(data, dict):
            return {key: self.sanitize_data(value) for key, value in data.items()}
        elif isinstance(data, list):
            return [self.sanitize_data(item) for item in data]
        elif isinstance(data, str):
            # 移除危险字符
            sanitized = data.replace('\x00', '')  # 移除null字节
            sanitized = re.sub(r'<script[^>]*>.*?</script>', '', sanitized, flags=re.IGNORECASE)
            return sanitized
        else:
            return data

    def log_security_event(self, event_type: str, details: dict):
        """记录安全事件"""
        log_data = {
            'timestamp': datetime.utcnow().isoformat(),
            'event_type': event_type,
            'client_ip': self.get_client_identifier(),
            'user_agent': request.headers.get('User-Agent', ''),
            'endpoint': request.endpoint,
            'method': request.method,
            'details': details
        }
        
        # 存储到Redis（可以后续改为专门的日志系统）
        log_key = f"security_log:{int(time.time())}"
        try:
            self.redis_client.setex(log_key, 86400 * 7, json.dumps(log_data))  # 保存7天
        except Exception as e:
            print(f"Failed to log security event: {str(e)}")

def rate_limit(endpoint: str):
    """请求频率限制装饰器"""
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if not hasattr(g, 'security_middleware'):
                return jsonify({
                    'success': False,
                    'message': '安全中间件未初始化',
                    'timestamp': datetime.utcnow().isoformat()
                }), 500
            
            # 检查频率限制
            allowed, info = g.security_middleware.check_rate_limit(endpoint)
            
            if not allowed:
                # 记录频率限制事件
                g.security_middleware.log_security_event('rate_limit_exceeded', {
                    'endpoint': endpoint,
                    'limit_info': info
                })
                
                response = jsonify({
                    'success': False,
                    'message': '请求过于频繁，请稍后再试',
                    'retry_after': info.get('retry_after', 60),
                    'timestamp': datetime.utcnow().isoformat()
                })
                response.status_code = 429
                response.headers['Retry-After'] = str(info.get('retry_after', 60))
                return response
            
            # 添加频率限制信息到响应头
            response = f(*args, **kwargs)
            if hasattr(response, 'headers'):
                response.headers['X-RateLimit-Remaining'] = str(info.get('remaining', 0))
                response.headers['X-RateLimit-Reset'] = str(info.get('reset_time', 0))
            
            return response
        return decorated_function
    return decorator

def validate_input():
    """输入验证装饰器"""
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if not hasattr(g, 'security_middleware'):
                return jsonify({
                    'success': False,
                    'message': '安全中间件未初始化',
                    'timestamp': datetime.utcnow().isoformat()
                }), 500
            
            # 验证请求数据
            if hasattr(g, 'decrypted_data') and g.decrypted_data:
                is_valid, message = g.security_middleware.validate_input_data(g.decrypted_data)
                
                if not is_valid:
                    # 记录安全事件
                    g.security_middleware.log_security_event('invalid_input', {
                        'validation_error': message,
                        'data_sample': str(g.decrypted_data)[:200]  # 只记录前200字符
                    })
                    
                    return jsonify({
                        'success': False,
                        'message': '输入数据验证失败',
                        'timestamp': datetime.utcnow().isoformat()
                    }), 400
                
                # 清理数据
                g.decrypted_data = g.security_middleware.sanitize_data(g.decrypted_data)
            
            return f(*args, **kwargs)
        return decorated_function
    return decorator

def get_security_headers() -> dict:
    """获取安全响应头"""
    return {
        'X-Content-Type-Options': 'nosniff',
        'X-Frame-Options': 'DENY',
        'X-XSS-Protection': '1; mode=block',
        'Strict-Transport-Security': 'max-age=31536000; includeSubDomains',
        'Content-Security-Policy': "default-src 'self'",
        'Referrer-Policy': 'strict-origin-when-cross-origin',
        'Permissions-Policy': 'geolocation=(), microphone=(), camera=()'
    }